# -*- coding: utf-8 -*-
import os
import matplotlib.pyplot as plt
from keras.applications.inception_v3 import InceptionV3
from keras.applications.resnet50 import ResNet50
from keras.applications.vgg16 import VGG16
from keras.models import Sequential
from keras.models import Model
from keras.utils import np_utils
from keras.layers import Convolution2D, Activation, MaxPooling2D, Flatten, Dense, Dropout, Input, Reshape, Lambda, \
    ZeroPadding2D
from keras import optimizers
from keras.preprocessing.image import ImageDataGenerator
from keras import backend as K
from keras.utils import plot_model
import numpy as np
from keras.callbacks import ModelCheckpoint, EarlyStopping, LearningRateScheduler, ReduceLROnPlateau


os.environ['CUDA_VISIBLE_DEVICES'] = '2'

def sign_sqrt(x):
    return K.sign(x) * K.sqrt(K.abs(x) + 1e-10)


def l2_norm(x):
    return K.l2_normalize(x, axis=-1)


def batch_dot(cnn_ab):
    print(cnn_ab[0].shape, cnn_ab[1].shape)
    # axes指定的轴作为列 行 然后必须相等，然后执行矩阵乘法
    return K.batch_dot(cnn_ab[0], cnn_ab[1], axes=[1, 1])

def vgg_reshape(x):
    _6layer = x.layers[-6].output
    p6layer = ZeroPadding2D(1)(_6layer)

    x = MaxPooling2D((2, 2), strides=(2, 2), name='block4_pool')(p6layer)
    # Block 5
    x = Convolution2D(512, (3, 3),
                      activation='relu',
                      padding='same',
                      name='block5_conv1')(x)
    x = Convolution2D(512, (3, 3),
                      activation='relu',
                      padding='same',
                      name='block5_conv2')(x)
    x = Convolution2D(512, (3, 3),
                      activation='relu',
                      padding='same',
                      name='block5_conv3')(x)
    x = MaxPooling2D((2, 2), strides=(2, 2), name='block5_pool')(x)
    # mdvg = Model(inputs=input_tensor,outputs=p6layer)
    # # mdvg = Model(inputs=p6layer,outputs=model_vgg16.layers[-5].output)
    # mdvgg = Model(inputs=mdvg,outputs=model_vgg16.output)
    # model_vg16 = Model(inputs=input_tensor,outputs=mdvgg)
    return x

def bilinearnet():
    input_tensor = Input(shape=(151, 72, 3))
    vgg16 = VGG16(weights='imagenet', include_top=False, input_tensor=input_tensor)
    #    conv2048 = Convolution2D(filters=2048,kernel_size=(3,3),)
    #    vgg16_add_conv_to_2048 = Model(inputs=input_tensor,outputs=)
    resnet50 = ResNet50(weights='imagenet', include_top=False, input_tensor=input_tensor)
    model_vgg16 = Model(inputs=input_tensor, outputs=vgg16.output)
    model_resnet50 = Model(inputs=input_tensor, outputs=resnet50.output)
    model_vgg16.compile(loss='categorical_crossentropy', optimizer='adam')
    model_resnet50.compile(loss='categorical_crossentropy', optimizer='adam')
    mdvg = vgg_reshape(model_vgg16)
    model_vgg16 = Model(inputs=input_tensor,outputs=mdvg)
    resnet50_x = Reshape([model_resnet50.layers[-35].output_shape[1] * model_resnet50.layers[-35].output_shape[2],
                          model_resnet50.layers[-35].output_shape[3]])(model_resnet50.layers[-35].output)
    vgg16_x = Reshape([model_vgg16.layers[-2].output_shape[1] * model_vgg16.layers[-2].output_shape[2],
                       model_vgg16.layers[-2].output_shape[3]])(model_vgg16.layers[-2].output)

    cnn_dot_out = Lambda(batch_dot)([vgg16_x, resnet50_x])

    sign_sqrt_out = Lambda(sign_sqrt)(cnn_dot_out)
    l2_norm_out = Lambda(l2_norm)(sign_sqrt_out)
    flatten = Flatten()(l2_norm_out)
    dropout = Dropout(0.5)(flatten)
    output = Dense(3, activation='softmax')(dropout)

    model = Model(input_tensor, output)
    model.compile(loss='categorical_crossentropy', optimizer=optimizers.SGD(lr=1e-4, momentum=0.9, decay=1e-6),
                  metrics=['accuracy'])
    print(model.summary())
    return model
    # plot_model(model, to_file='vgg_resnet_bilinear_model.png')


# input_tensor = Input(shape=(151, 72, 3))
# vgg16 = VGG16(weights='imagenet',include_top=False, input_tensor=input_tensor)
# #    conv2048 = Convolution2D(filters=2048,kernel_size=(3,3),)
# #    vgg16_add_conv_to_2048 = Model(inputs=input_tensor,outputs=)
# resnet50 = ResNet50(weights='imagenet', include_top=False, input_tensor=input_tensor)
# model_vgg16 = Model(inputs=input_tensor, outputs=vgg16.output)
# model_resnet50 = Model(inputs=input_tensor, outputs=resnet50.output)
#
#
# _6layer = model_vgg16.layers[-6].output
# p6layer = ZeroPadding2D(1)(_6layer)
# x = MaxPooling2D((2, 2), strides=(2, 2), name='block4_pool')(p6layer)
# # Block 5
# x = Convolution2D(512, (3, 3),
#                   activation='relu',
#                   padding='same',
#                   name='block5_conv1')(x)
# x = Convolution2D(512, (3, 3),
#                   activation='relu',
#                   padding='same',
#                   name='block5_conv2')(x)
# x = Convolution2D(512, (3, 3),
#                   activation='relu',
#                   padding='same',
#                   name='block5_conv3')(x)
# x = MaxPooling2D((2, 2), strides=(2, 2), name='block5_pool')(x)
# # mdvg = Model(inputs=input_tensor,outputs=p6layer)
# # # mdvg = Model(inputs=p6layer,outputs=model_vgg16.layers[-5].output)
# # mdvgg = Model(inputs=mdvg,outputs=model_vgg16.output)
# # model_vg16 = Model(inputs=input_tensor,outputs=mdvgg)
# mdvg = Model(inputs=input_tensor,outputs=x)
# print(f'--------',model_resnet50.layers[-35].output)
# print(f'vgg16\n', mdvg.layers[-2].output)
# print(f'\nresnet', model_resnet50.summary())

train_data_gen = ImageDataGenerator(rescale=1 / 255.,
                                    samplewise_center=True,
                                    samplewise_std_normalization=True,
                                    #                 zca_whitening=True,
                                    #                 zca_epsilon=1e-6,

                                    width_shift_range=0.05,
                                    height_shift_range=0.05,
                                    fill_mode='reflect',
                                    horizontal_flip=True,
                                    vertical_flip=True)

test_data_gen = ImageDataGenerator(rescale=1 / 255.)
#
train_gen = train_data_gen.flow_from_directory(directory='/data/soft/redldw/shuang',
                                               target_size=(151,72), color_mode='rgb',
                                               class_mode='categorical',
                                               batch_size=5, shuffle=True, seed=222
                                               )

val_gen = test_data_gen.flow_from_directory(directory='/data/soft/redldw/shuang',
                                            target_size=(151,72), color_mode='rgb',
                                            class_mode='categorical',
                                            batch_size=5, shuffle=True, seed=2
                                            )
test_gen = test_data_gen.flow_from_directory(directory='/data/soft/redldw/shuang',
                                             target_size=(151,72), color_mode='rgb',
                                             class_mode='categorical',
                                             batch_size=5,seed=100,
                                             )
cp = ModelCheckpoint('guangdong_best_vgg16.h5', monitor='val_loss', verbose=1,
                     save_best_only=True, save_weights_only=False,
                     mode='auto', period=1)
es = EarlyStopping(monitor='val_loss',
                   patience=8, verbose=1, mode='auto')
lr_reduce = ReduceLROnPlateau(monitor='val_loss', factor=0.1, epsilon=1e-5, patience=2, verbose=1, min_lr=0.00000001)
callbackslist = [es, lr_reduce]

ftvggmodel = bilinearnet()
ftvggmodel.fit_generator(train_gen,
                         epochs=3,
                         verbose=1,
                         steps_per_epoch=1915,
                         callbacks=callbackslist,
                         validation_data=val_gen,
                         validation_steps=1915,
                         shuffle=True)

# ftvggmodel.load_weights('guangdong_best_vgg16.h5')
pred = ftvggmodel.predict_generator(test_gen,)
ftvggmodel.save("./bishuang.h5")
defectlist = ['good','bad','other']
import csv

with open('lvcai_result.csv', 'w') as f:
    w = csv.writer(f)
    for i in range(len(pred)):
        w.writerow([str(i) + '.jpg', defectlist[np.argmax(pred[i])]])
